import os
from pathlib import Path
import cv2
import numpy as np
import matplotlib.pyplot as plt
import torch
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from utils import pred_lines,pred_squares
from models.mbv2_mlsd_large import MobileV2_MLSD_Large

class LineDetecter:
    def __init__(self, mlsd_model_path):
        
        # 加载MLSD模型
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.mlsd_model = MobileV2_MLSD_Large().to(self.device).eval()
        self.mlsd_model.load_state_dict(
            torch.load(mlsd_model_path, map_location=self.device, weights_only=True), 
            strict=True
        )
    
    def detect_lines_left_bottom(self,img,left_bottom,line_threshold=0.1, dist_threshold=20):
        """检测left_bottom顶点附近的直线"""
        # 以left_bottom为中心裁剪img为126x126的图像
        h, w = img.shape[:2]
        x_start = int(max(0, left_bottom[0] - 63))
        y_start = int(max(0, left_bottom[1] - 63))
        cropped_img = img[y_start:y_start+126, x_start:x_start+126]
        
        # 将裁剪图像resize到512x512用于MLSD
        # cropped_resized = cv2.resize(cropped_img, (512, 512))
        
        # 转换颜色空间
        cropped_rgb = cv2.cvtColor(cropped_img, cv2.COLOR_BGR2RGB)
        
        # 使用MLSD检测直线
        lines = pred_lines(cropped_rgb, self.mlsd_model, [126, 126], line_threshold, dist_threshold)
        
        # 将直线坐标转换回裁剪图像尺寸
        scale_x = 126 / 126
        scale_y = 126 / 126
        lines_cropped = []
        for line in lines:
            x1, y1, x2, y2 = line
            x1_crop = int(x1 * scale_x)
            y1_crop = int(y1 * scale_y)
            x2_crop = int(x2 * scale_x)
            y2_crop = int(y2 * scale_y)
            lines_cropped.append([x1_crop, y1_crop, x2_crop, y2_crop])
        
        # 将裁剪图像中的直线坐标转换回原图尺寸
        lines_original = []
        for line in lines_cropped:
            x1_crop, y1_crop, x2_crop, y2_crop = line
            x1_orig = x1_crop + x_start
            y1_orig = y1_crop + y_start
            x2_orig = x2_crop + x_start
            y2_orig = y2_crop + y_start
            lines_original.append([x1_orig, y1_orig, x2_orig, y2_orig])
        
        return lines_original

    def detect_lines(self, img, line_threshold=0.1, dist_threshold=20):
        """检测图像中的直线"""
        # 将图像resize到512x512用于MLSD
        h, w = img.shape[:2]
        img_resized = cv2.resize(img, (512, 512))
        
        # 转换颜色空间
        img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)
        
        # 使用MLSD检测直线
        lines = pred_lines2(img_rgb, self.mlsd_model, [512, 512], line_threshold, dist_threshold)
        
        # 将直线坐标转换回原图尺寸
        scale_x = w / 512
        scale_y = h / 512
        lines_original = []
        
        for line in lines:
            x1, y1, x2, y2 = line
            x1_orig = int(x1 * scale_x)
            y1_orig = int(y1 * scale_y)
            x2_orig = int(x2 * scale_x)
            y2_orig = int(y2 * scale_y)
            lines_original.append([x1_orig, y1_orig, x2_orig, y2_orig])
        
        return lines_original
